package com.cyf.common.core.coll;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

/**
 * 带有最大容量限制的优先队列，如果队列已满，并且新元素大于队列中的最小元素，则删除最小元素并添加新元素。
 * @author chenyifan
 * @create 2024-09-23 15:02
 */
public class MaxNHeap<E> {

    private PriorityQueue<E> queue;
    private Comparator<? super E> comparator;
    // 堆的大小
    private int maxSize;

    public MaxNHeap(int maxSize) {
        this(maxSize, null);
    }

    public MaxNHeap(int maxSize, Comparator<? super E> comparator) {
        this.maxSize = maxSize;
        this.comparator = comparator;
        // 小顶堆
        this.queue = new PriorityQueue<>(maxSize, comparator);
    }

    public void add(E e) {
        if (queue.size() < maxSize) {
            queue.offer(e);
        } else if (compare(e, queue.peek())){
            // 新元素大于队首元素，删除队首元素，添加新元素
            queue.poll();
            queue.offer(e);
        }
    }

    public List<E> getList() {
        return new ArrayList<>(queue);
    }

    /**
     * 比较 e1, e2 大小
     * @param e1
     * @param e2
     * @return e1 > e2 返回 true，否则返回 false
     */
    private boolean compare(E e1, E e2) {
        if (comparator != null) {
            return comparator.compare(e1, e2) > 0;
        } else {
            return ((Comparable<? super E>) e1).compareTo(e2) > 0;
        }
    }

}
